package com.fanshuai.net.proxy.impl;

import com.fanshuai.net.proxy.Proxy;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.http.*;
import io.netty.handler.codec.http.websocketx.*;
import io.netty.handler.stream.ChunkedWriteHandler;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import java.io.IOException;
import java.net.URI;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
public class ProxyClient implements Proxy {
    private final String host;
    private final int port;
    private final SocketChannel outBoundChannel; //客户端socket连接

    private EventLoopGroup httpGroup;
    private EventLoopGroup webSocketGroup;

    private Channel httpChannel;
    private Channel webSocketChannel;
    private WebSocketClientHandshaker clientHandShaker;

    private final AtomicBoolean httpInit = new AtomicBoolean(false);
    private final AtomicBoolean webSocketInit = new AtomicBoolean(false);

    static final String DEFAULT_HOST = "127.0.0.1";
    static final int DEFAULT_PORT = 80;

    public ProxyClient(String serverAddr, SocketChannel socketChannel) {
        if (StringUtils.isBlank(serverAddr) || !serverAddr.contains(":")) {
            host = DEFAULT_HOST;
            port = DEFAULT_PORT;
        } else {
            String[] arr = serverAddr.split(":");
            host = arr[0];
            port = Integer.parseInt(arr[1]);
        }

        this.outBoundChannel = socketChannel;
    }

    public void initHttpClient() {
        if (httpInit.get()) {
            return;
        }

        httpGroup = new NioEventLoopGroup(1);

        Bootstrap httpBootstrap = new Bootstrap();
        httpBootstrap.group(httpGroup).channel(NioSocketChannel.class).handler(new ChannelInitializer<SocketChannel>() {
            @Override
            protected void initChannel(SocketChannel ch) throws Exception {
                ChannelPipeline pipeline = ch.pipeline();
                pipeline.addLast(new HttpClientCodec(), new HttpObjectAggregator(65536), new ChunkedWriteHandler());

                ClientHttpHandler httpHandler = new ClientHttpHandler(outBoundChannel);
                pipeline.addLast(httpHandler);
            }
        });

        try {
            ChannelFuture f = httpBootstrap.connect(host, port).sync();
            httpChannel = f.channel();
            log.info("proxy http connect success, channel={}", f.channel());

            httpInit.set(true);
        } catch (Exception e) {
            log.error("proxy http connect error, ", e);
        }
    }

    public void webSocketHandshake(FullHttpRequest request) {
        if (webSocketInit.get()) {
            return;
        }

        webSocketGroup = new NioEventLoopGroup(1);

        Bootstrap webSocketBootstrap = new Bootstrap();

        String path = request.uri();
        HttpHeaders headers = request.headers();
        //webSocket handshaker
        URI uri = URI.create(String.format("ws://%s:%d%s", host, port, path));
        this.clientHandShaker = WebSocketClientHandshakerFactory.newHandshaker(uri, WebSocketVersion.V13, null, false, headers);

        ClientWebSocketHandler webSocketHandler = new ClientWebSocketHandler(outBoundChannel, clientHandShaker);
        webSocketBootstrap.group(webSocketGroup).channel(NioSocketChannel.class).handler(new ChannelInitializer<SocketChannel>() {
            @Override
            protected void initChannel(SocketChannel ch) throws Exception {
                ChannelPipeline pipeline = ch.pipeline();
                pipeline.addLast(new HttpClientCodec(),
                        new HttpObjectAggregator(65536),
                        new ChunkedWriteHandler());

                pipeline.addLast(webSocketHandler);
            }
        });

        try {
            //tcp连接
            ChannelFuture f = webSocketBootstrap.connect(host, port).sync();
            webSocketChannel = f.channel();
            log.info("proxy webSocket connect success, channel={}", f.channel());

            log.info("proxy websocket begin handshake, uri={}, headers={}", uri, getHeaders(headers));
            webSocketHandler.getHandshakeFuture().sync();

            webSocketInit.set(true);
        } catch (Exception e) {
            log.error("proxy websocket connect error, ", e);
        }
    }

    //将http请求发送到代理服务
    public void sendHttpRequest(FullHttpRequest copyRequest) {
        if (!httpInit.get()) {
            initHttpClient();
        }

        httpChannel.writeAndFlush(copyRequest);
    }

    //将webSocket消息转发到代理服务
    public void sendWebSocketFrame(WebSocketFrame copyFrame) {
        if (!webSocketInit.get()) {
            throw new RuntimeException("websocket is not init");
        }

        webSocketChannel.writeAndFlush(copyFrame);
    }

    @Override
    public void close() throws IOException {
        if (httpInit.get()) {
            httpGroup.shutdownGracefully();
            httpChannel.close();
        }
        if (webSocketInit.get()) {
            webSocketGroup.shutdownGracefully();

            //webSocket连接关闭
            clientHandShaker.close(webSocketChannel, new CloseWebSocketFrame());
            webSocketChannel.close();
        }
    }

    private static String getHeaders(HttpHeaders headers) {
        StringBuilder b = new StringBuilder();
        b.append("\n");
        for (Map.Entry<String, String> entry : headers.entries()) {
            b.append(String.format("%s: %s\n", entry.getKey(), entry.getValue()));
        }
        return b.toString();
    }

    static class ClientHttpHandler extends SimpleChannelInboundHandler<Object> {
        private final SocketChannel outBoundChannel;
        public ClientHttpHandler(SocketChannel outBoundChannel) {
            this.outBoundChannel = outBoundChannel;
        }

        @Override
        public void channelActive(ChannelHandlerContext ctx) throws Exception {
            super.channelActive(ctx);
            log.info("proxy http channelActive, channel={}", ctx.channel());
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            log.info("proxy http channelInactive, channel={}", ctx.channel());
            super.channelInactive(ctx);
        }

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, Object msg) {
            if (!(msg instanceof FullHttpResponse)) {
                throw new RuntimeException("unexpected frame, msg=" + msg);
            }
            try {
                FullHttpResponse response = (FullHttpResponse) msg;
                log.info("receive proxy http response, header={}", getHeaders(response.headers()));
                outBoundChannel.writeAndFlush(response.copy());
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    static class ClientWebSocketHandler extends SimpleChannelInboundHandler<Object> {
        private final SocketChannel outBoundChannel;
        private final WebSocketClientHandshaker clientHandshaker;
        private ChannelPromise handshakeFuture;

        public ClientWebSocketHandler(SocketChannel outBoundChannel, WebSocketClientHandshaker handshaker) {
            this.outBoundChannel = outBoundChannel;
            this.clientHandshaker = handshaker;
        }

        public ChannelPromise getHandshakeFuture() {
            return handshakeFuture;
        }

        @Override
        public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
            super.handlerAdded(ctx);
            handshakeFuture = ctx.newPromise();
        }

        @Override
        public void channelActive(ChannelHandlerContext ctx) throws Exception {
            super.channelActive(ctx);
            log.info("proxy webSocket channelActive, channel={}", ctx.channel());
            clientHandshaker.handshake(ctx.channel());
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            log.info("proxy webSocket channelInactive, channel={}", ctx.channel());
            super.channelInactive(ctx);
        }

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, Object msg) {
            if (!clientHandshaker.isHandshakeComplete()) {
                try {
                    FullHttpResponse response = (FullHttpResponse) msg;
                    clientHandshaker.finishHandshake(ctx.channel(), response);
                    log.info("proxy websocket handshake success");
                    handshakeFuture.setSuccess();
                } catch (Exception e) {
                    log.error("proxy websocket handshake failed, ", e);
                    handshakeFuture.setFailure(e);
                }
            } else {
                if (msg instanceof CloseWebSocketFrame) {
                    log.info("close webSocket");
                    ctx.channel().close();
                }
                if (msg instanceof PongWebSocketFrame) {
                    log.info("frame={}", msg);
                }
                if (msg instanceof TextWebSocketFrame) {
                    TextWebSocketFrame frame = (TextWebSocketFrame) msg;
                    outBoundChannel.writeAndFlush(frame.copy());
                    log.info("receive proxy websocket response, content={}", frame.content());
                }
            }
        }
    }
}
